{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Change Detection (Segmentation) using TAO Visual ChangeNet-Segmentation Network\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## What is a Visual ChangeNet Network?\n",
    "\n",
    "Visual ChangeNet is a state of the art transformer-based Change Detection model. Similar to the optical inspection notebook, Visual ChangeNet is also based on Siamese Network, which is a class of neural network architectures containing two or more identical subnetworks. The training algorithm works by updating the parameters across all the sub-networks in tandem. In TAO, Visual ChangeNet supports two images as input where the end goal is to either classify or segment the change between the \"golden or reference\" image and the \"test\" image. TAO supports the [FAN](https://arxiv.org/abs/2204.12451) backbone network for both Visual ChangeNet architectures. For more details about training FAN backbones, please refer to the Classification pytorch notebook.\n",
    "\n",
    "TAO Toolkit versions 5.3 and later support some of the foundational models for change detection (classification and segmentation). NV-DINOv2 can now be used as the backbone for the Visual ChangeNet-Classification and Segmentation models.\n",
    "\n",
    "To mitigate the inferiror performance of a standard vision transformer (ViT) on dense prediction tasks, TAO supports the ViT-Adapter architecture. This allows a powerful ViT that has learned rich semantic representations from a large corpus of data to achieve comparable performance to vision-specific transformers on dense preidiction tasks.\n",
    "\n",
    "In TAO, two different types of Change Detection networks are supported: Visual ChangeNet-Segmentation and Visual ChangeNet-Classification intended for Segmentation and Classification of change between the two input images, respectively. Visual ChangeNet-Segmentation is specifically intended for segmentation change detection. In this notebook, we leverage the Visual ChangeNet-Segmentation model to demonstrate change detection by segmenting the changed pixels between the two input images. \n",
    "\n",
    "### Sample inference from the Visual ChangeNet-Segmentation model\n",
    "\n",
    "Here is a sample inference output for the Visual ChangeNet-Segmentation model. The model takes in two input images (test and golden reference) and outputs a segmentation change map for building change detection on LEVIR-CD dataset.\n",
    "\n",
    "<img align=\"center\" title=\"no defects\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/purpose_built_models/visual_changenet/sample_image.jpg?raw=true\" width=\"800\" >\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Train a Visual ChangeNet-Segmentation Model for Semantic change detection on remote sensing imagery.\n",
    "* Evaluate the trained model.\n",
    "* Run Inference on the trained model.\n",
    "* Export the trained model to a .onnx file (encrypted ONNX model) for deployment to DeepStream or TensorRT.\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `visual_changenet` model, \n",
    "which you may deploy with this [end-to-end sample](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps) with Triton.\n",
    "\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of Siamese Network using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset](#head-2)\n",
    "3. [Download the pretrained model from NGC](#head-3)\n",
    "4. [Provide training specification](#head-4)\n",
    "5. [Run TAO training](#head-5)\n",
    "6. [Evaluate trained models](#head-6)\n",
    "7. [Inferences](#head-7)\n",
    "8. [Deploy](#head-8)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=FIXME\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"changenet\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"changenet\", \"results\")\n",
    "os.environ[\"HOST_MODEL_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"changenet\", \"model\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/changenet\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR\n",
    "! mkdir -p $HOST_MODEL_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_MODEL_DIR\"],\n",
    "           \"destination\": \"/model\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.7.0. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 7 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7.0 < 3.8.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 525+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.7.0. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare a sample dataset <a class=\"anchor\" id=\"head-2\"></a>\n",
    "\n",
    "TAO Visual ChangeNet-Segmentation uses a custom dataset format. The sections below walk through this format, specifically how the dataset should be structured and the files required."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visual ChangeNet-Segmentation data format\n",
    "\n",
    "Visual ChangeNet-Segmentation expects directories of images and mask files in the dataset root directory. The image directories consist of\n",
    "golden image directory (pre-change images) and test image directory (post-change image) to compare against the change mask images with pixel level change masks.  \n",
    "\n",
    "\n",
    "     |--dataset_root:\n",
    "          |--A\n",
    "               |--image1.jpg\n",
    "               |--image2.jpg\n",
    "          |--B\n",
    "               |--image1.jpg\n",
    "               |--image2.jpg\n",
    "          |--label\n",
    "               |--image1.jpg\n",
    "               |--image2.jpg      \n",
    "          |--list\n",
    "               |-- train.txt\n",
    "               |-- val.txt\n",
    "               |-- test.txt\n",
    "               |-- predict.txt\n",
    "\n",
    "Here's a description of the structure:\n",
    "\n",
    "* The ``dataset_root`` directory contains the following:\n",
    "    * ``A``: Contains post-change test images\n",
    "    * ``B``: Contains pre-change golden reference images\n",
    "    * ``label``: Contains ground truth segmentation change masks\n",
    "    * ``list``: Contains .txt files for each dataset split, as described in the  section below. \n",
    "\n",
    "#### List Files\n",
    "\n",
    "\n",
    "Visual ChangeNet-Segmentation dataloader expects the ``label`` directory to contain ``.txt`` files for each of the dataset split [train, validation, test, inference]. \n",
    "A Visual ChangeNet-Segmentation label file is a simple ``.txt`` file containing all file names for the particular split.\n",
    "\n",
    "\n",
    "| ``image_names``   | \n",
    "| :---------------- |\n",
    "| ``file_name.png`` |\n",
    "\n",
    "* ``image_names``: The names of images. Image names should be the same for test images and their corresponding reference and mask images. \n",
    "\n",
    "Here is a sample label file corresponding to the sample directory structure as describe in the section above.\n",
    "\n",
    "\n",
    "| ``image_names``   | \n",
    "| :---------------- |\n",
    "| ``image1.png``    |\n",
    "| ``image3.png``    |\n",
    "| ``image2.png``    |\n",
    "\n",
    "Note: Each test image (inside directory ``A``) must have a reference image (inside directory ``B``)\n",
    "and a segmentation change map (inside directory ``label``) with the same name for the dataloader to map them correctly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Download the data\n",
    "import os\n",
    "LEVIR_DOWNLOAD_URL = \"https://www.dropbox.com/s/18fb5jo0npu5evm/LEVIR-CD256.zip\"\n",
    "os.environ[\"URL_DATASET\"]=LEVIR_DOWNLOAD_URL\n",
    "!if [ ! -f $HOST_DATA_DIR/LEVIR-CD256.zip ]; then wget $URL_DATASET -O $HOST_DATA_DIR/LEVIR-CD-256.zip; else echo \"image archive already downloaded\"; fi "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Verify the downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!mkdir -p $HOST_DATA_DIR\n",
    "!if [ ! -f $HOST_DATA_DIR/LEVIR-CD-256.zip ]; then echo 'Dataset zip file not found, please download.'; else echo 'Found Dataset zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!unzip -u $HOST_DATA_DIR/LEVIR-CD-256.zip -d $HOST_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "def verify_dataset(data_dir):\n",
    "\n",
    "    num_images_A = len(os.listdir(os.path.join(DATA_DIR, \"LEVIR-CD256/A\")))\n",
    "    num_images_B = len(os.listdir(os.path.join(DATA_DIR, \"LEVIR-CD256/B\")))\n",
    "    num_labels = len(os.listdir(os.path.join(DATA_DIR, \"LEVIR-CD256/label\")))\n",
    "    num_train_images = sum(1 for _ in open(os.path.join(DATA_DIR, \"LEVIR-CD256/list/train.txt\")))\n",
    "    num_val_images = sum(1 for _ in open(os.path.join(DATA_DIR, \"LEVIR-CD256/list/val.txt\")))\n",
    "    num_test_images = sum(1 for _ in open(os.path.join(DATA_DIR, \"LEVIR-CD256/list/test.txt\")))\n",
    "    print(\"Number of images in the train/val/test set. {}\".format(num_images_A))\n",
    "    print(\"Number of compare images in the train/val/test set. {}\".format(num_images_B))\n",
    "    print(\"Number of labels in the train/val/test set. {}\".format(num_labels))\n",
    "    print(\"Number of train instances. {}\".format(num_train_images))\n",
    "    print(\"Number of val instances. {}\".format(num_val_images))\n",
    "    print(\"Number of test instances. {}\".format(num_test_images))\n",
    "\n",
    "    # # Use test set for inference\n",
    "    # shutil.copy(os.path.join(DATA_DIR, \"LEVIR-CD-256/list/test.txt\"), os.path.join(DATA_DIR, \"LEVIR-CD-256/list/predict.txt\"))\n",
    "\n",
    "DATA_DIR = os.environ.get('HOST_DATA_DIR')\n",
    "# verify downloaded dataset\n",
    "verify_dataset(data_dir=DATA_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Verify\n",
    "!ls -l $HOST_DATA_DIR/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Download the pretrained model from NGC <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "Download the pretrained model from NGC. We will use NGC CLI to get the data and model. For more details, go to https://ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "import os\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# # Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip\n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/visual_changenet_segmentation_levircd:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model download-version \"nvidia/tao/visual_changenet_segmentation_levircd:visual_changenet_levircd_trainable_v1.0\" --dest $HOST_RESULTS_DIR/pretrained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained/visual_changenet_segmentation_levircd_vvisual_changenet_levircd_trainable_v1.0/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Provide training specification <a class=\"anchor\" id=\"head-4\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* task: segment\n",
    "* train:\n",
    "  * resume_training_checkpoint_path: null\n",
    "  * pretrained_model_path: /path/to/pretrained/model.pth\n",
    "  * segment:\n",
    "    * loss: \"ce\"\n",
    "    * weights: [0.5, 0.5, 0.5, 0.8, 1.0]\n",
    "  * num_epochs: 350\n",
    "  * num_nodes: 1\n",
    "  * validation_interval: 1\n",
    "  * checkpoint_interval: 1\n",
    "  * optim:\n",
    "    * lr: 0.0002\n",
    "    * optim: \"adamw\"\n",
    "    * policy: \"linear\" \n",
    "    * momentum: 0.9\n",
    "    * weight_decay: 0.01\n",
    "    * betas: [0.9, 0.999]\n",
    "* results_dir: \"/results\"\n",
    "* model:\n",
    "  * backbone:\n",
    "    * type: \"fan_small_12_p4_hybrid\"\n",
    "    * pretrained_backbone_path: /path/to/pretrained/backbone.pth\n",
    "    * freeze_backbone: False\n",
    "  * decode_head:\n",
    "    * feature_strides: [4, 8, 16, 16]\n",
    "* dataset:\n",
    "  * segment:\n",
    "    * dataset: \"CNDataset\"\n",
    "    * root_dir: /data/LEVIR-CD-256\n",
    "    * data_name: \"LEVIR-CD\"\n",
    "    * label_transform: \"norm\"\n",
    "    * batch_size: 16\n",
    "    * workers: 2\n",
    "    * multi_scale_train: True\n",
    "    * multi_scale_infer: False\n",
    "    * num_classes: 2\n",
    "    * img_size: 256\n",
    "    * image_folder_name: \"A\"\n",
    "    * change_image_folder_name: 'B'\n",
    "    * list_folder_name: 'list'\n",
    "    * annotation_folder_name: \"label\"\n",
    "    * train_split: \"train\"\n",
    "    * validation_split: \"val\"\n",
    "    * test_split: 'test'\n",
    "    * predict_split: 'predict'\n",
    "    * label_suffix: .png\n",
    "    * augmentation: \n",
    "      * random_flip:\n",
    "        * vflip_probability: 0.5\n",
    "        * hflip_probability: 0.5\n",
    "        * enable: True\n",
    "      * random_rotate:\n",
    "        * rotate_probability: 0.5\n",
    "        * angle_list: [90, 180, 270]\n",
    "        * enable: True\n",
    "      * random_color:\n",
    "        * brightness: 0.3\n",
    "        * contrast: 0.3\n",
    "        * saturation: 0.3\n",
    "        * hue: 0.3\n",
    "        * enable: True\n",
    "      * with_scale_random_crop:\n",
    "        * enable: True\n",
    "      * with_random_crop: True\n",
    "      * with_random_blur: True\n",
    "\n",
    "Please refer to the TAO documentation about Visual ChangeNet-Segmentation to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/experiment.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Run TAO training <a class=\"anchor\" id=\"head-5\"></a>\n",
    "* Provide the sample spec file and the output directory location for models.\n",
    "* WARNING: Training will take several hours or one day to complete.\n",
    "* LEVIR-CD per-epoch training time on a single GPU (the hours may vary depending on the data location, network speed, and etc).\n",
    "\n",
    "<table>\n",
    "  <tr>\n",
    "    <th>GPU Type</th>\n",
    "    <th>Time (minutes)</th>\n",
    "  </tr>\n",
    "  <tr>\n",
    "    <td>1 x A6000 48GB</td>\n",
    "    <td>12</td>\n",
    "  </tr>\n",
    "</table>\n",
    "\n",
    "* If you're looking to speed through the notebook, you can adjust the `train.num_epochs` parameter for a quicker run. It's crucial to understand that this adjustment is intended for demonstration and educational purposes only and won't produce state-of-the-art (SOTA) results.\n",
    "* To replicate baseline numbers accurately, the model should be trained for 300 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env MODEL_DIR = /model\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results\n",
    "%env NUM_EPOCHS=5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Train Visual ChangeNet-Segmentation model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will train a Visual ChangeNet-Segmentation model from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Train model\")\n",
    "!tao model visual_changenet train \\\n",
    "                  -e $SPECS_DIR/experiment.yaml \\\n",
    "                    train.pretrained_model_path=$RESULTS_DIR/pretrained/visual_changenet_segmentation_levircd_vvisual_changenet_levircd_trainable_v1.0/changenet_segment_levir_cd.pth\n",
    "                    # train.num_epochs=$NUM_EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training command for multi-gpu training. We can define the number of gpus using the `train.num_gpus` parameter.\n",
    "## The following command will trigger multi-gpu training on 2 gpus.\n",
    "# print(\"Train model\")\n",
    "# !tao model visual_changenet train \\\n",
    "#                   -e $SPECS_DIR/experiment.yaml \\\n",
    "#                   train.num_gpus=2\n",
    "#                   train.num_epochs=$NUM_EPOCHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model checkpoints:')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/changenet_model_segment_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/changenet_segment.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/changenet_segment.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Evaluate trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "Evaluate trained model.\n",
    "\n",
    "For Visual ChangeNet-Segmentation model evaluation, the following metrics are preferred:\n",
    "\n",
    "* Overall Accuracy.\n",
    "* mIoU: Intersection over Union (IoU) for all classes.\n",
    "\n",
    "Note that we provide metrics for Precision, Recall, IoU and F1 score for each individual class as well as averaged over all classes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model visual_changenet evaluate \\\n",
    "                   -e $SPECS_DIR/experiment.yaml \\\n",
    "                    evaluate.checkpoint=$RESULTS_DIR/train/changenet_segment.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Visualize Inference <a class=\"anchor\" id=\"head-7\"></a>\n",
    "In this section, we run the Visual ChangeNet-Segmentation inference tool to generate inferences with the trained models and save the results under `$RESULTS_DIR`. We also visualise a few sample outputs along with their corresponding input images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model visual_changenet inference \\\n",
    "                   -e $SPECS_DIR/experiment.yaml \\\n",
    "                    inference.checkpoint=$RESULTS_DIR/train/changenet_segment.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images(output_path, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"inference\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Deploy <a class=\"anchor\" id=\"head-8\"></a>\n",
    "Export the model to encrypted ONNX model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model visual_changenet export \\\n",
    "                    -e $SPECS_DIR/experiment.yaml \\\n",
    "                        export.checkpoint=$RESULTS_DIR/train/changenet_segment.pth \\\n",
    "                        export.onnx_file=$RESULTS_DIR/export/changenet.onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Uncomment this cell to export an onnx file with dynamic batching enabled for\n",
    "# # integration with trtexec and deepstream.\n",
    "# # Export the model with export.batch_size=-1 for dynamic batching.\n",
    "# ! tao model visual_changenet export \\\n",
    "#                     -e $SPECS_DIR/experiment.yaml \\\n",
    "#                     export.checkpoint=$RESULTS_DIR/train/changenet_segment.pth \\\n",
    "#                     export.batch_size=-1 \\\n",
    "#                     export.results_dir=$RESULTS_DIR/export_dynamic \\\n",
    "#                     export.onnx_file=$RESULTS_DIR/export_dynamic/changenet.onnx\n",
    "                            \n",
    "\n",
    "# # Profiling the exported model via trtexec.\n",
    "# ! tao deploy trtexec --onnx=$RESULTS_DIR/export_dynamic/changenet.onnx \\\n",
    "#                      --minShapes=input0:1x3x512x128,input1:1x3x256x256 \\\n",
    "#                      --optShapes=input0:8x3x512x128,input1:8x3x256x256 \\\n",
    "#                      --maxShapes=input0:16x3x512x128,input1:16x3x256x256 \\\n",
    "#                      --fp16 \\\n",
    "#                      --saveEngine=$RESULTS_DIR/export_dynamic/changenet-fp16.engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lth $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet gen_trt_engine \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "                gen_trt_engine.onnx_file=$RESULTS_DIR/export/changenet.onnx \\\n",
    "                gen_trt_engine.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet.trt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet inference \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "                inference.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet.trt \\\n",
    "                inference.results_dir=$RESULTS_DIR/trt_inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy visual_changenet evaluate \\\n",
    "            -e $SPECS_DIR/experiment.yaml \\\n",
    "                evaluate.trt_engine=$RESULTS_DIR/gen_trt_engine/changenet.trt \\\n",
    "                evaluate.results_dir=$RESULTS_DIR/trt_evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"trt_inference\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
